{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.3"
    },
    "colab": {
      "name": "Robotic Manipulation - Basic Pick and Place.ipynb",
      "provenance": [],
      "collapsed_sections": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EgiF12Hf1Dhs",
        "colab_type": "text"
      },
      "source": [
        "**I recommend you run the first code cell of this notebook immediately, to start provisioning drake on the cloud machine, then you can leave this window open as you [read the textbook](manipulation.csail.mit.edu/pick.html).**\n",
        "\n",
        "# Notebook setup\n",
        "\n",
        "The following cell will:\n",
        "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
        "- launch a server for our 3D visualizer (MeshCat) that will be used for the remainder of this notebook.\n",
        "\n",
        "You will need to rerun this cell if you restart the kernel, but it should be fast because the machine will already have drake installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eeMrMI0-1Dhu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import importlib\n",
        "import sys\n",
        "from urllib.request import urlretrieve\n",
        "\n",
        "# Install drake.\n",
        "if 'google.colab' in sys.modules and importlib.util.find_spec('pydrake') is None:\n",
        "  version='20200918'\n",
        "  build='nightly'\n",
        "  urlretrieve(f\"https://drake-packages.csail.mit.edu/drake/{build}/drake-{version}/setup_drake_colab.py\",\n",
        "              \"setup_drake_colab.py\")\n",
        "  from setup_drake_colab import setup_drake\n",
        "  setup_drake(version=version, build=build)\n",
        "  !pip install pyngrok==4.2.2\n",
        "\n",
        "# Determine if this notebook is currently running as a notebook or a unit test.\n",
        "from IPython import get_ipython\n",
        "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
        "\n",
        "# Install pyngrok.\n",
        "server_args = []\n",
        "if 'google.colab' in sys.modules:\n",
        "  server_args = ['--ngrok_http_tunnel']\n",
        "\n",
        "# Start a single meshcat server instance to use for the remainder of this notebook.\n",
        "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
        "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
        "\n",
        "\n",
        "# Let's do all of our imports here, too.\n",
        "import numpy as np\n",
        "import altair as alt\n",
        "import pydot\n",
        "from IPython.display import display, SVG\n",
        "from ipywidgets import Text, Textarea, Layout\n",
        "import plotly.express as px\n",
        "\n",
        "from pydrake.all import (\n",
        "    AddMultibodyPlantSceneGraph, AngleAxis, BasicVector, ConnectMeshcatVisualizer, \n",
        "    DiagramBuilder, FindResourceOrThrow, Integrator, JacobianWrtVariable, \n",
        "    LeafSystem, MultibodyPlant, MultibodyPositionToGeometryPose, Parser, \n",
        "    PiecewisePolynomial, PiecewiseQuaternionSlerp, Quaternion, RigidTransform, \n",
        "    RollPitchYaw, RotationMatrix, SceneGraph, Simulator, TrajectorySource\n",
        ")\n",
        "from pydrake.examples.manipulation_station import ManipulationStation\n",
        "from pydrake.multibody.jupyter_widgets import MakeJointSlidersThatPublishOnCallback\n",
        "\n",
        "# TODO(russt): Move this to drake (adding the element name support to the base class).\n",
        "import pandas as pd\n",
        "\n",
        "def dataframe(trajectory, times, names):\n",
        "  assert trajectory.rows() == len(names)\n",
        "  values = trajectory.vector_values(times)\n",
        "  data = {'t': times }\n",
        "  for i in range(len(names)):\n",
        "    data[names[i]] = values[i,:]\n",
        "  return pd.DataFrame(data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7q0A14bAilIX",
        "colab_type": "text"
      },
      "source": [
        "# Inspecting the kinematic tree\n",
        "\n",
        "Here is a simple example that demonstrates how to inspect the kinematic tree stored in a `MultibodyPlant`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ILYLouFTjv6e",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def kinematic_tree_example():\n",
        "    plant = MultibodyPlant(time_step=0.0)\n",
        "    parser = Parser(plant)\n",
        "    parser.AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/manipulation/models/allegro_hand_description/sdf/allegro_hand_description_right.sdf\"))\n",
        "    parser.AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\"))\n",
        "    plant.Finalize()\n",
        "\n",
        "    # TODO(russt): Add floating base connections\n",
        "    # TODO(russt): Consider a more interactive javascript rendering?\n",
        "    display(SVG(pydot.graph_from_dot_data(plant.GetTopologyGraphvizString())[0].create_svg()))\n",
        "\n",
        "kinematic_tree_example()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wSa5QcU0Piak",
        "colab_type": "text"
      },
      "source": [
        "# Forward kinematics of the gripper frame\n",
        "\n",
        "Here is a simple example that let's you visualize the frames on the iiwa and the gripper.  If you click on the \"Open Controls\" menu in the MeshCat visualizer, and dig into the menu `meshcat->drake->Source` then you will see elements for each of the models in the `SceneGraph`: one for the iiwa, another for the WSG, and others for the clutter bins.  You can enable/disable their visualization.  Give it a spin!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5SjOClhTltPk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def gripper_forward_kinematics_example():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    # TODO: Replace this with a simple model directive of iiwa+wsg (no clutter bins)\n",
        "    station = builder.AddSystem(ManipulationStation())\n",
        "    station.SetupClutterClearingStation()\n",
        "    station.Finalize()\n",
        "\n",
        "    frames_to_draw = {\"iiwa\": {\"iiwa_link_1\", \"iiwa_link_2\", \"iiwa_link_3\", \"iiwa_link_4\", \"iiwa_link_5\", \"iiwa_link_6\", \"iiwa_link_7\"},\n",
        "                      \"gripper\": {\"body\"}}\n",
        "    meshcat = ConnectMeshcatVisualizer(builder,\n",
        "        station.get_scene_graph(),\n",
        "        output_port=station.GetOutputPort(\"pose_bundle\"),\n",
        "        zmq_url=zmq_url,\n",
        "        frames_to_draw=frames_to_draw,\n",
        "        axis_length=0.3,\n",
        "        axis_radius=0.01)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    xyz = Text(value=\"\", description=\"gripper position (m): \", layout=Layout(width='500px'), style={'description_width':'initial'})\n",
        "    rpy = Text(value=\"\", description=\"gripper roll-pitch-yaw (rad): \", layout=Layout(width='500px'), style={'description_width':'initial'})\n",
        "    plant = station.get_multibody_plant()\n",
        "\n",
        "    gripper = plant.GetBodyByName(\"body\")\n",
        "    def pose_callback(context):\n",
        "        pose = plant.EvalBodyPoseInWorld(context, gripper)   ## This is the important line\n",
        "        xyz.value = np.array2string(pose.translation(), formatter={'float': lambda x: \"{:3.2f}\".format(x)})\n",
        "        rpy.value = np.array2string(RollPitchYaw(pose.rotation()).vector(), formatter={'float': lambda x: \"{:3.2f}\".format(x)})\n",
        "\n",
        "    meshcat.load()\n",
        "    MakeJointSlidersThatPublishOnCallback(station.get_multibody_plant(), meshcat, context, my_callback=pose_callback)\n",
        "    display(xyz)\n",
        "    display(rpy)\n",
        "\n",
        "gripper_forward_kinematics_example()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ZUhwdTOtOjp",
        "colab_type": "text"
      },
      "source": [
        "# Don't assume $\\dot{q} \\equiv v$\n",
        "\n",
        "Let's just add a single object into the scene.  We won't weld it to the world frame, so it is a \"free body\" or has a \"floating base\"."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MKezK0PjtUP4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def num_positions_velocities_example():\n",
        "    plant = MultibodyPlant(time_step = 0.0)\n",
        "    Parser(plant).AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\"))\n",
        "    plant.Finalize()\n",
        "\n",
        "    context = plant.CreateDefaultContext()\n",
        "    print(context)\n",
        "\n",
        "    print(f\"plant.num_positions() = {plant.num_positions()}\")\n",
        "    print(f\"plant.num_velocities() = {plant.num_velocities()}\")\n",
        "\n",
        "num_positions_velocities_example()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tlE6mkO0ajJu",
        "colab_type": "text"
      },
      "source": [
        "Looking at the `Context` you can see that this system has 13 total state variables.  7 of them are positions, $q$; this is due to our pose representation using unit quaternions.  But only 6 of them are velocities, $v$; this is because a six-element spatial velocity provides a better (unconstrained) representation of the rate of change of the unit quaternion.  But clearly, if the length of the vectors don't even match, we do *not* have $\\dot{q} = v$."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hSeD3PBotJUU",
        "colab_type": "text"
      },
      "source": [
        "# Kinematic Jacobians for pick and place\n",
        "\n",
        "Let's set up the same iiwa + wsg example, with sliders (but without the frames), that we used above.  But this time I'll display the value of the Jacobian $J^G(q)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6F_KZoc1QDkY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def pick_and_place_jacobians_example():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    # TODO: Replace this with a simple model directive of iiwa+wsg (no clutter bins)\n",
        "    station = builder.AddSystem(ManipulationStation())\n",
        "    station.SetupClutterClearingStation()\n",
        "    station.Finalize()\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(builder,\n",
        "        station.get_scene_graph(),\n",
        "        output_port=station.GetOutputPort(\"pose_bundle\"),\n",
        "#        delete_prefix_on_load=False,\n",
        "        zmq_url=zmq_url)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    jacobian = Textarea(value=\"\", description=\"J_G: \", layout={'width':'60%','height':'100px'}, style={'description_width':'initial'})\n",
        "    sigma = Text(value=\"\", description=\"smallest singular value(J_G): \", layout={'width':'60%'}, style={'description_width':'initial'})\n",
        "    plant = station.get_multibody_plant()\n",
        "\n",
        "    G = plant.GetBodyByName(\"body\").body_frame()\n",
        "    W = plant.world_frame()\n",
        "    def pose_callback(context):\n",
        "        J_G = plant.CalcJacobianSpatialVelocity(context, JacobianWrtVariable.kQDot, G, [0,0,0], W, W)   ## This is the important line\n",
        "        jacobian.value = np.array2string(J_G, formatter={'float': lambda x: \"{:5.1f}\".format(x)})\n",
        "        sigma.value = str(np.min(np.linalg.svd(J_G, compute_uv=False)))\n",
        "\n",
        "    meshcat.load()\n",
        "\n",
        "    # If you want to set the initial positions manually, use this:\n",
        "    #plant.SetPositions(plant.GetMyContextFromRoot(context),\n",
        "    #                   plant.GetModelInstanceByName(\"iiwa\"),\n",
        "    #                   [0, 0, 0, 0, 0, 0, 0])\n",
        "\n",
        "    MakeJointSlidersThatPublishOnCallback(station.get_multibody_plant(), meshcat, context, my_callback=pose_callback)\n",
        "    display(jacobian)\n",
        "    display(sigma)\n",
        "\n",
        "pick_and_place_jacobians_example()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZym7Xo7slPU",
        "colab_type": "text"
      },
      "source": [
        "# Our first end-effector \"controller\"\n",
        "\n",
        "Let's use the pseudo-inverse of the Jacobian to drive the robot around.  To do that, we'll write a very simple system that looks at the current value of $q$, computes $[J^G]^+$, and uses it to command a constant spatial velocity, $V^G$.\n",
        "\n",
        "We'll only run this controller for a short duration.  Constant spatial velocities aren't something that makes sense for a longer simulation!\n",
        "\n",
        "Make sure you try changing $V^G$, and understand how the command relates to the motion of the robot."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WBTFzB5ZuXvP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# We can write a new System by deriving from the LeafSystem class.\n",
        "# There is a little bit of boiler plate, but hopefully this example makes sense.\n",
        "class PseudoInverseController(LeafSystem):\n",
        "    def __init__(self, plant):\n",
        "        LeafSystem.__init__(self)\n",
        "        self._plant = plant\n",
        "        self._plant_context = plant.CreateDefaultContext()\n",
        "        self._iiwa = plant.GetModelInstanceByName(\"iiwa\")\n",
        "        self._G = plant.GetBodyByName(\"body\").body_frame()\n",
        "        self._W = plant.world_frame()\n",
        "\n",
        "        self.DeclareVectorInputPort(\"iiwa_position\", BasicVector(7))\n",
        "        self.DeclareVectorOutputPort(\"iiwa_velocity\", BasicVector(7), \n",
        "                                     self.CalcOutput)\n",
        "\n",
        "    def CalcOutput(self, context, output):\n",
        "        q = self.get_input_port().Eval(context)\n",
        "        self._plant.SetPositions(self._plant_context, self._iiwa, q)\n",
        "        J_G = self._plant.CalcJacobianSpatialVelocity(\n",
        "            self._plant_context, JacobianWrtVariable.kQDot, \n",
        "            self._G, [0,0,0], self._W, self._W)\n",
        "        J_G = J_G[:,0:7] # Ignore gripper terms\n",
        "        \n",
        "        V_G_desired = np.array([0,    # rotation about x\n",
        "                                -.1,  # rotation about y\n",
        "                                0,    # rotation about z\n",
        "                                0,    # x\n",
        "                                -.05, # y\n",
        "                                -.1]) # z\n",
        "        v = np.linalg.pinv(J_G).dot(V_G_desired)\n",
        "        output.SetFromVector(v)\n",
        "\n",
        "        \n",
        "def jacobian_controller_example():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    station = builder.AddSystem(ManipulationStation())\n",
        "    station.SetupClutterClearingStation()\n",
        "    station.Finalize()\n",
        "\n",
        "    controller = builder.AddSystem(PseudoInverseController(\n",
        "        station.get_multibody_plant()))\n",
        "    integrator = builder.AddSystem(Integrator(7))\n",
        "\n",
        "    builder.Connect(controller.get_output_port(), \n",
        "                    integrator.get_input_port())\n",
        "    builder.Connect(integrator.get_output_port(),\n",
        "                    station.GetInputPort(\"iiwa_position\"))\n",
        "    builder.Connect(station.GetOutputPort(\"iiwa_position_measured\"),\n",
        "                    controller.get_input_port())\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(builder,\n",
        "        station.get_scene_graph(),\n",
        "        output_port=station.GetOutputPort(\"pose_bundle\"),\n",
        "        zmq_url=zmq_url)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    simulator = Simulator(diagram)\n",
        "    station_context = station.GetMyContextFromRoot(simulator.get_mutable_context())\n",
        "    station.GetInputPort(\"iiwa_feedforward_torque\").FixValue(station_context, np.zeros((7,1)))\n",
        "    station.GetInputPort(\"wsg_position\").FixValue(station_context, [0.1])\n",
        "    # TODO(russt): Add this missing python binding\n",
        "    #integrator.set_integral_value(\n",
        "    #    integrator.GetMyContextFromRoot(simulator.get_mutable_context()), \n",
        "    #        station.GetIiwaPosition(station_context))\n",
        "    integrator.GetMyContextFromRoot(simulator.get_mutable_context()).get_mutable_continuous_state_vector().SetFromVector(station.GetIiwaPosition(station_context))\n",
        "\n",
        "    simulator.set_target_realtime_rate(1.0)\n",
        "    simulator.AdvanceTo(0.01);\n",
        "\n",
        "    return simulator\n",
        "\n",
        "simulator = jacobian_controller_example()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6v-EGfoI3y6V",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Wait for the model to load in the visualizer, then run this cell to see the interesting part...\n",
        "simulator.AdvanceTo(5.0 if running_as_notebook else 0.1);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DfwIMIJCSkx6",
        "colab_type": "text"
      },
      "source": [
        "# Compute grasp and pregrasp poses\n",
        "\n",
        "Here is a simple example with a floating Schunk gripper and the foam brick.  It defines the grasp pose as described in the notes, and renders it to the 3D visualizer.\n",
        "\n",
        "**Check yourself**: Try changing the grasp pose to our pregrasp pose.  Do you like the numbers that I picked in the text?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QuFmomDOxoAR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def grasp_poses_example():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step = 0.0)\n",
        "    parser = Parser(plant, scene_graph)\n",
        "    grasp = parser.AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf\"), \"grasp\")\n",
        "    # TODO(russt): Draw the pregrasp gripper, too, as transparent (drake #13970).\n",
        "    #pregrasp = parser.AddModelFromFile(FindResourceOrThrow(\n",
        "    #    \"drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf\"), \"pregrasp\")\n",
        "    brick = parser.AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\"))\n",
        "    plant.Finalize()\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "    plant_context = plant.GetMyContextFromRoot(context)\n",
        "\n",
        "    # TODO(russt): Set a random pose of the object.\n",
        "\n",
        "    # Get the current object, O, pose\n",
        "    B_O = plant.GetBodyByName(\"base_link\", brick)\n",
        "    X_WO = plant.EvalBodyPoseInWorld(plant_context, B_O)\n",
        "\n",
        "    B_Ggrasp = plant.GetBodyByName(\"body\", grasp)\n",
        "    p_GgraspO = [0, 0.12, 0]\n",
        "    R_GgraspO = RotationMatrix.MakeXRotation(np.pi/2.0).multiply(\n",
        "        RotationMatrix.MakeZRotation(np.pi/2.0))\n",
        "    X_GgraspO = RigidTransform(R_GgraspO, p_GgraspO)\n",
        "    X_OGgrasp = X_GgraspO.inverse()\n",
        "    X_WGgrasp = X_WO.multiply(X_OGgrasp)\n",
        "\n",
        "    plant.SetFreeBodyPose(plant_context, B_Ggrasp, X_WGgrasp)\n",
        "    # Open the fingers, too.\n",
        "    plant.GetJointByName(\"left_finger_sliding_joint\", grasp).set_translation(plant_context, -0.054)\n",
        "    plant.GetJointByName(\"right_finger_sliding_joint\", grasp).set_translation(plant_context, 0.054)\n",
        "\n",
        "    meshcat.load()\n",
        "    diagram.Publish(context)\n",
        "\n",
        "grasp_poses_example()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d5p2tbJWrWZ5",
        "colab_type": "text"
      },
      "source": [
        "# A pick and place \"sketch\" with keyframes\n",
        "\n",
        "Let's first define the canonical frames of our problem.  We've got a handful of them, and will pass them in and out of several functions, so I'll use a dictionary to store them.\n",
        "\n",
        "*Edit*: I've added one more frame into the mix here, to get some clearance when moving between the bins.  Leaving it out gave my a nice [robot outtake](http://manipulation.csail.mit.edu/data/pick_fail.html), but we can do better."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zYuRDF6mSBx5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_gripper_frames(X_G, X_O):\n",
        "  \"\"\"\n",
        "  Takes a partial specification with X_G[\"initial\"] and X_O[\"initial\"] and X_0[\"goal\"], and \n",
        "  returns a X_G and times with all of the pick and place frames populated.\n",
        "  \"\"\"\n",
        "  # Define (again) the gripper pose relative to the object when in grasp.\n",
        "  p_GgraspO = [0, 0.12, 0]\n",
        "  R_GgraspO = RotationMatrix.MakeXRotation(np.pi/2.0).multiply(\n",
        "      RotationMatrix.MakeZRotation(np.pi/2.0))\n",
        "  X_GgraspO = RigidTransform(R_GgraspO, p_GgraspO)\n",
        "  X_OGgrasp = X_GgraspO.inverse()\n",
        "  # pregrasp is negative y in the gripper frame (see the figure!).\n",
        "  X_GgraspGpregrasp = RigidTransform([0, -0.08, 0])\n",
        "\n",
        "  X_G[\"pick\"] = X_O[\"initial\"].multiply(X_OGgrasp)\n",
        "  X_G[\"prepick\"] = X_G[\"pick\"].multiply(X_GgraspGpregrasp)\n",
        "  X_G[\"place\"] = X_O[\"goal\"].multiply(X_OGgrasp)\n",
        "  X_G[\"preplace\"] = X_G[\"place\"].multiply(X_GgraspGpregrasp)\n",
        "\n",
        "  # I'll interpolate a ahlfway orientation by converting to axis angle and halving the angle.\n",
        "  X_GprepickGpreplace = X_G[\"prepick\"].inverse().multiply(X_G[\"preplace\"])\n",
        "  angle_axis = X_GprepickGpreplace.rotation().ToAngleAxis()\n",
        "  X_GprepickGclearance = RigidTransform(AngleAxis(angle=angle_axis.angle()/2.0, axis=angle_axis.axis()), \n",
        "                                        X_GprepickGpreplace.translation()/2.0 + np.array([0, -0.3, 0]))\n",
        "  X_G[\"clearance\"] = X_G[\"prepick\"].multiply(X_GprepickGclearance)\n",
        "\n",
        "  # Now let's set the timing\n",
        "  times = {\"initial\": 0}\n",
        "  X_GinitialGprepick = X_G[\"initial\"].inverse().multiply(X_G[\"prepick\"])\n",
        "  times[\"prepick\"] = times[\"initial\"] + 10.0*np.linalg.norm(X_GinitialGprepick.translation())\n",
        "  # Allow some time for the gripper to close.\n",
        "  times[\"pick_start\"] = times[\"prepick\"] + 2.0\n",
        "  times[\"pick_end\"] = times[\"pick_start\"] + 2.0\n",
        "  times[\"postpick\"] = times[\"pick_end\"] + 2.0\n",
        "  time_to_from_clearance = 10.0*np.linalg.norm(X_GprepickGclearance.translation())\n",
        "  times[\"clearance\"] = times[\"postpick\"] + time_to_from_clearance\n",
        "  times[\"preplace\"] = times[\"clearance\"] + time_to_from_clearance\n",
        "  times[\"place_start\"] = times[\"preplace\"] + 2.0\n",
        "  times[\"place_end\"] = times[\"place_start\"] + 2.0\n",
        "  times[\"postplace\"] = times[\"place_end\"] + 2.0\n",
        "\n",
        "  return X_G, times\n",
        "\n",
        "X_G = {\"initial\": RigidTransform(RotationMatrix.MakeXRotation(-np.pi/2.0), [0, -0.25, 0.25])}\n",
        "X_O = {\"initial\": RigidTransform(RotationMatrix.MakeZRotation(np.pi/2.0), [-.2, -.75, 0.025]),\n",
        "       \"goal\": RigidTransform(RotationMatrix.MakeZRotation(np.pi),[.75, 0, 0.025])}\n",
        "X_G, times = make_gripper_frames(X_G, X_O)\n",
        "print(f\"Sanity check: The entire maneuver will take {times['postplace']} seconds to execute.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JnmfEw0zgzsn",
        "colab_type": "text"
      },
      "source": [
        "Here's a simple visualization of those frames.  Use the controls menu in meshcat to view the different objects on the different \"paths\" to understand which objects belong to which frames.  Everything look reasonable?\n",
        "\n",
        "Try changing the initial and final pose of the gripper / object."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oXh50Cf1hU6N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def visualize_gripper_frames(X_G, X_O):\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step = 0.0)\n",
        "    parser = Parser(plant, scene_graph)\n",
        "    gripper = FindResourceOrThrow(\n",
        "        \"drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf\")\n",
        "    brick = FindResourceOrThrow(\n",
        "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\")\n",
        "    for key, pose in X_G.items():\n",
        "      g = parser.AddModelFromFile(gripper, f\"gripper_{key}\")\n",
        "      plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"body\", g), pose)\n",
        "    for key, pose in X_O.items():\n",
        "      o = parser.AddModelFromFile(brick, f\"object_{key}\")\n",
        "      plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"base_link\", o), pose)\n",
        "\n",
        "    plant.Finalize()\n",
        "\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url)\n",
        "    meshcat.load()\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "    diagram.Publish(context)\n",
        "\n",
        "visualize_gripper_frames(X_G, X_O)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HmcG3b_xSeQR",
        "colab_type": "text"
      },
      "source": [
        "# Pick and place trajectories\n",
        "\n",
        "We've defined the keyframes, but now we will fill in timing details (what should be happening at each instant).  Interpolating between the positions is quite simple; we can encode these details as a `PiecewisePolynomial` trajectory in $\\Re^3$ over time.\n",
        "\n",
        "TODO(russt): Maybe implement `PiecewisePose` with an append workflow."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oiwljuszWh_O",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_gripper_position_trajectory(X_G, times):\n",
        "    \"\"\"\n",
        "    Constructs a gripper position trajectory from the plan \"sketch\".\n",
        "    \"\"\"\n",
        "    # The syntax here is a little ugly; we need to pass in a matrix with the samples in the columns.\n",
        "    # TODO(russt): Add python bindings for the std::vector constructor.\n",
        "    traj = PiecewisePolynomial.FirstOrderHold(\n",
        "        [times[\"initial\"], times[\"prepick\"]], np.vstack([X_G[\"initial\"].translation(), X_G[\"prepick\"].translation()]).T)\n",
        "\n",
        "    # TODO(russt): I could make this less brittle if I was more careful on the names above, and just look up the pose for every time (in order)\n",
        "    traj.AppendFirstOrderSegment(times[\"pick_start\"], X_G[\"pick\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"pick_end\"], X_G[\"pick\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"postpick\"], X_G[\"prepick\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"clearance\"], X_G[\"clearance\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"preplace\"], X_G[\"preplace\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"place_start\"], X_G[\"place\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"place_end\"], X_G[\"place\"].translation())\n",
        "    traj.AppendFirstOrderSegment(times[\"postplace\"], X_G[\"preplace\"].translation())\n",
        "\n",
        "    return traj\n",
        "\n",
        "traj_p_G = make_gripper_position_trajectory(X_G, times)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2PnYrnSmEMYC",
        "colab_type": "text"
      },
      "source": [
        "Now let's plot our trajectories.  `Altair` is a plotting library built on Pandas and Vega-lite (to support interactive graphics), which works well in Colab.  I aim to start using it more.  \n",
        "\n",
        "<sup><sub>TODO(russt): Add mouseover interaction that publishes to meshcat, so that a cursor on this plot renders the corresponding pose.(Unfortunately, I don't believe I can link a callback to python with altair).</sub></sup>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ppd-5B3V3cNC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = dataframe(traj_p_G, traj_p_G.get_segment_times(), ['x','y','z'])\n",
        "alt.vconcat(\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='x').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='y').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='z').properties(height=80),\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yIVvA6ZoFLvt",
        "colab_type": "text"
      },
      "source": [
        "Unfortunately, Altair doesn't do 3D plotting (neither does mpld3), so here is the Plotly Express version."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2vgn7mxcDvlK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = dataframe(traj_p_G, np.linspace(traj_p_G.start_time(), traj_p_G.end_time(), 50), ['x','y','z'])\n",
        "fig = px.line_3d(data, x='x', y='y', z='z')\n",
        "fig.update_traces(mode='lines+markers',\n",
        "                  marker=dict(size=3, color=\"rgb(153, 51, 51)\"),\n",
        "                  line=dict(color=\"rgb(153, 51, 51)\"))\n",
        "if running_as_notebook:\n",
        "    fig.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NhXx70MP4C2N",
        "colab_type": "text"
      },
      "source": [
        "We also need to come up with a time trajectory for the orientations of our gripper.  Simple straight-line interpolation between rotations doesn't work particularly well, no matter how you are representing your rotations.  We'll use the quaternion \"slerp\" here (short for \"spherical linear interpolation\"), using the `PiecewiseQuaternionSlerp` trajectory class."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4yvktzauCntQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_gripper_orientation_trajectory(X_G, times):\n",
        "    \"\"\"\n",
        "    Constructs a gripper position trajectory from the plan \"sketch\".\n",
        "    \"\"\"\n",
        "    traj = PiecewiseQuaternionSlerp();\n",
        "    traj.Append(times[\"initial\"], X_G[\"initial\"].rotation())\n",
        "    traj.Append(times[\"prepick\"], X_G[\"prepick\"].rotation())\n",
        "    traj.Append(times[\"pick_start\"], X_G[\"pick\"].rotation())\n",
        "    traj.Append(times[\"pick_end\"], X_G[\"pick\"].rotation())\n",
        "    traj.Append(times[\"postpick\"], X_G[\"prepick\"].rotation())\n",
        "    traj.Append(times[\"clearance\"], X_G[\"clearance\"].rotation())\n",
        "    traj.Append(times[\"preplace\"], X_G[\"preplace\"].rotation())\n",
        "    traj.Append(times[\"place_start\"], X_G[\"place\"].rotation())\n",
        "    traj.Append(times[\"place_end\"], X_G[\"place\"].rotation())\n",
        "    traj.Append(times[\"postplace\"], X_G[\"preplace\"].rotation())\n",
        "\n",
        "    return traj\n",
        "\n",
        "traj_R_G = make_gripper_orientation_trajectory(X_G, times)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zNNRr4Fq7frG",
        "colab_type": "text"
      },
      "source": [
        "We can plot the orientation trajectory, too.  Here are the components of the quaternion over time."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4aKEnaLJ50pF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = dataframe(traj_R_G, traj_R_G.get_segment_times(), ['x','y','z','w'])\n",
        "alt.vconcat(\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='x').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='y').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='z').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='w').properties(height=80)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWCyowMC8LUN",
        "colab_type": "text"
      },
      "source": [
        "Better, I think, to just plot the entire result.  Before we do, let's quickly define the finger trajectory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "65OocrXD8TSA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "opened = np.array([0.107]);\n",
        "closed = np.array([0.0]);\n",
        "\n",
        "def make_wsg_command_trajectory(times):\n",
        "    traj_wsg_command = PiecewisePolynomial.FirstOrderHold(\n",
        "        [times[\"initial\"], times[\"pick_start\"]], np.hstack([[opened], [opened]]))\n",
        "    traj_wsg_command.AppendFirstOrderSegment(times[\"pick_end\"], closed)\n",
        "    traj_wsg_command.AppendFirstOrderSegment(times[\"place_start\"], closed)\n",
        "    traj_wsg_command.AppendFirstOrderSegment(times[\"place_end\"], opened)\n",
        "    traj_wsg_command.AppendFirstOrderSegment(times[\"postplace\"], opened)\n",
        "    return traj_wsg_command\n",
        "\n",
        "traj_wsg_command = make_wsg_command_trajectory(times)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aEc95stq9tyx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = dataframe(traj_wsg_command, traj_wsg_command.get_segment_times(), ['wsg'])\n",
        "alt.Chart(data).mark_line().encode(x='t', y='wsg').properties(height=80)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wrhTvm5eIJFi",
        "colab_type": "text"
      },
      "source": [
        "Now let's really visualize the trajectory.  The best way is to animate it in 3D!  Please forgive the messiness of this particular bit of code; I'm trying to provide some intermediate steps that aren't typically used by the main drake developers.\n",
        "\n",
        "TODO(russt): make this less \"bespoke\"."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PZ7S7oai_ORF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class GripperTrajectoriesToPosition(LeafSystem):\n",
        "    def __init__(self, plant, traj_p_G, traj_R_G, traj_wsg_command):\n",
        "        LeafSystem.__init__(self)\n",
        "        self.plant = plant\n",
        "        self.gripper_body = plant.GetBodyByName(\"body\")\n",
        "        self.left_finger_joint = plant.GetJointByName(\"left_finger_sliding_joint\")\n",
        "        self.right_finger_joint = plant.GetJointByName(\"right_finger_sliding_joint\")\n",
        "        self.traj_p_G = traj_p_G\n",
        "        self.traj_R_G = traj_R_G\n",
        "        self.traj_wsg_command = traj_wsg_command\n",
        "        self.plant_context = plant.CreateDefaultContext()\n",
        "\n",
        "        self.DeclareVectorOutputPort(\"position\", \n",
        "                                     BasicVector(plant.num_positions()),\n",
        "                                     self.CalcPositionOutput)\n",
        "\n",
        "    def CalcPositionOutput(self, context, output):\n",
        "        t = context.get_time()\n",
        "        X_G = RigidTransform(Quaternion(self.traj_R_G.value(t)), self.traj_p_G.value(t))\n",
        "        self.plant.SetFreeBodyPose(self.plant_context, self.gripper_body, X_G)\n",
        "        wsg = self.traj_wsg_command.value(t)\n",
        "        self.left_finger_joint.set_translation(self.plant_context, -wsg/2.0)\n",
        "        self.right_finger_joint.set_translation(self.plant_context, wsg/2.0)\n",
        "        output.SetFromVector(self.plant.GetPositions(self.plant_context))\n",
        "\n",
        "def visualize_pick_and_place_trajectory(traj_p_G, traj_R_G, traj_wsg_command, X_O):\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    # Note: Don't use AddMultibodyPlantSceneGraph because we are only using\n",
        "    # MultibodyPlant for parsing, then wiring directly to SceneGraph.\n",
        "    scene_graph = builder.AddSystem(SceneGraph())\n",
        "    plant = MultibodyPlant(time_step=0.0)\n",
        "    plant.RegisterAsSourceForSceneGraph(scene_graph)\n",
        "    parser = Parser(plant, scene_graph)\n",
        "    gripper = parser.AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/manipulation/models/wsg_50_description/sdf/schunk_wsg_50_no_tip.sdf\"), \"gripper\")\n",
        "    brick = FindResourceOrThrow(\n",
        "        \"drake/examples/manipulation_station/models/061_foam_brick.sdf\")\n",
        "    for key, pose in X_O.items():\n",
        "      o = parser.AddModelFromFile(brick, f\"object_{key}\")\n",
        "      plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"base_link\", o), pose)\n",
        "    plant.Finalize()\n",
        "\n",
        "    to_pose = builder.AddSystem(MultibodyPositionToGeometryPose(plant))\n",
        "    builder.Connect(to_pose.get_output_port(), scene_graph.get_source_pose_port(plant.get_source_id()))\n",
        "\n",
        "    traj_to_position = builder.AddSystem(GripperTrajectoriesToPosition(plant, traj_p_G, traj_R_G, traj_wsg_command))\n",
        "    builder.Connect(traj_to_position.get_output_port(), to_pose.get_input_port())\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "\n",
        "    simulator = Simulator(diagram)\n",
        "    meshcat.start_recording()\n",
        "    simulator.AdvanceTo(traj_p_G.end_time() if running_as_notebook else 0.1)\n",
        "    meshcat.publish_recording()\n",
        "\n",
        "visualize_pick_and_place_trajectory(traj_p_G, traj_R_G, traj_wsg_command, X_O)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ue9ofS7GHpXr",
        "colab_type": "text"
      },
      "source": [
        "# Putting it all together\n",
        "\n",
        "In the visualization above, we animated the plan, but did not actually command the robot!  Let's put our grasping trajectories back together with our Jacobian-based controller.\n",
        "\n",
        "The first thing you might have noticed is that the Jacobian controller takes spatial velocities as input.  But so far we have trajectories for positions and orientations.  Let's differentiate them to get the velocities."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MfaCoTnNXfjj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "traj_v_G = traj_p_G.MakeDerivative()\n",
        "traj_w_G = traj_R_G.MakeDerivative()\n",
        "\n",
        "data = dataframe(traj_v_G, np.linspace(traj_v_G.start_time(), traj_v_G.end_time()), ['x','y','z'])\n",
        "alt.vconcat(\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='x').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='y').properties(height=80),\n",
        "  alt.Chart(data).mark_line().encode(x='t', y='z').properties(height=80),\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BrDNL8cvMo_o",
        "colab_type": "text"
      },
      "source": [
        "Now we need a slightly modified version of our Jacobian controller, which can take the spatial velocity commands from an input port."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nVKlTQKGKUU0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# We can write a new System by deriving from the LeafSystem class.\n",
        "# There is a little bit of boiler plate, but hopefully this example makes sense.\n",
        "class PseudoInverseController(LeafSystem):\n",
        "    def __init__(self, plant):\n",
        "        LeafSystem.__init__(self)\n",
        "        self._plant = plant\n",
        "        self._plant_context = plant.CreateDefaultContext()\n",
        "        self._iiwa = plant.GetModelInstanceByName(\"iiwa\")\n",
        "        self._G = plant.GetBodyByName(\"body\").body_frame()\n",
        "        self._W = plant.world_frame()\n",
        "\n",
        "        self.w_G_port = self.DeclareVectorInputPort(\"omega_WG\", BasicVector(3))\n",
        "        self.v_G_port = self.DeclareVectorInputPort(\"v_WG\", BasicVector(3))\n",
        "        self.q_port = self.DeclareVectorInputPort(\"iiwa_position\", BasicVector(7))\n",
        "        self.DeclareVectorOutputPort(\"iiwa_velocity\", BasicVector(7), \n",
        "                                     self.CalcOutput)\n",
        "        # TODO(russt): Add missing binding\n",
        "        #joint_indices = plant.GetJointIndices(self._iiwa)\n",
        "        #self.position_start = plant.get_joint(joint_indices[0]).position_start()\n",
        "        #self.position_end = plant.get_joint(joint_indices[-1]).position_start()\n",
        "        self.iiwa_start = plant.GetJointByName(\"iiwa_joint_1\").velocity_start()\n",
        "        self.iiwa_end = plant.GetJointByName(\"iiwa_joint_7\").velocity_start()\n",
        "\n",
        "    def CalcOutput(self, context, output):\n",
        "        w_G = self.w_G_port.Eval(context)\n",
        "        v_G = self.v_G_port.Eval(context)\n",
        "        V_G = np.hstack([w_G, v_G])\n",
        "        q = self.q_port.Eval(context)\n",
        "        self._plant.SetPositions(self._plant_context, self._iiwa, q)\n",
        "        J_G = self._plant.CalcJacobianSpatialVelocity(\n",
        "            self._plant_context, JacobianWrtVariable.kV, \n",
        "            self._G, [0,0,0], self._W, self._W)\n",
        "        J_G = J_G[:,self.iiwa_start:self.iiwa_end+1] # Only iiwa terms.\n",
        "        v = np.linalg.pinv(J_G).dot(V_G)\n",
        "        output.SetFromVector(v)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uM1pk1zzN9D3",
        "colab_type": "text"
      },
      "source": [
        "Oh, boy.  Big moment.  I'm excited... are you?  Let's do this thing!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hDrXl1oSNt6u",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "X_O = {\"initial\": RigidTransform(RotationMatrix.MakeZRotation(np.pi/2.0), [-.2, -.65, 0.09]),\n",
        "       \"goal\": RigidTransform(RotationMatrix.MakeZRotation(np.pi),[.5, 0, 0.09])}\n",
        "\n",
        "builder = DiagramBuilder()\n",
        "\n",
        "station = builder.AddSystem(ManipulationStation())\n",
        "station.SetupClutterClearingStation()\n",
        "station.AddManipulandFromFile(\n",
        "    \"drake/examples/manipulation_station/models/061_foam_brick.sdf\",\n",
        "    X_O[\"initial\"])\n",
        "station.Finalize()\n",
        "\n",
        "# Find the initial pose of the gripper (as set in the default Context)\n",
        "temp_context = station.CreateDefaultContext()\n",
        "plant = station.get_multibody_plant()\n",
        "temp_plant_context = plant.GetMyContextFromRoot(temp_context)\n",
        "X_G = {\"initial\": plant.EvalBodyPoseInWorld(temp_plant_context, plant.GetBodyByName(\"body\"))}\n",
        "X_G, times = make_gripper_frames(X_G, X_O)\n",
        "print(f\"Sanity check: The entire maneuver will take {times['postplace']} seconds to execute.\")\n",
        "\n",
        "# Make the trajectories\n",
        "traj_p_G = make_gripper_position_trajectory(X_G, times)\n",
        "traj_v_G = traj_p_G.MakeDerivative()\n",
        "traj_R_G = make_gripper_orientation_trajectory(X_G, times)\n",
        "traj_w_G = traj_R_G.MakeDerivative()\n",
        "\n",
        "v_G_source = builder.AddSystem(TrajectorySource(traj_v_G))\n",
        "v_G_source.set_name(\"v_WG\")\n",
        "w_G_source = builder.AddSystem(TrajectorySource(traj_w_G))\n",
        "w_G_source.set_name(\"omega_WG\")\n",
        "controller = builder.AddSystem(PseudoInverseController(plant))\n",
        "controller.set_name(\"PseudoInverseController\")\n",
        "builder.Connect(v_G_source.get_output_port(), controller.GetInputPort(\"v_WG\"))\n",
        "builder.Connect(w_G_source.get_output_port(), controller.GetInputPort(\"omega_WG\"))\n",
        "\n",
        "integrator = builder.AddSystem(Integrator(7))\n",
        "integrator.set_name(\"integrator\")\n",
        "builder.Connect(controller.get_output_port(), \n",
        "                integrator.get_input_port())\n",
        "builder.Connect(integrator.get_output_port(),\n",
        "                station.GetInputPort(\"iiwa_position\"))\n",
        "builder.Connect(station.GetOutputPort(\"iiwa_position_measured\"),\n",
        "                controller.GetInputPort(\"iiwa_position\"))\n",
        "\n",
        "traj_wsg_command = make_wsg_command_trajectory(times)\n",
        "wsg_source = builder.AddSystem(TrajectorySource(traj_wsg_command))\n",
        "wsg_source.set_name(\"wsg_command\")\n",
        "builder.Connect(wsg_source.get_output_port(), station.GetInputPort(\"wsg_position\"))\n",
        "\n",
        "meshcat = ConnectMeshcatVisualizer(builder,\n",
        "    station.get_scene_graph(),\n",
        "    output_port=station.GetOutputPort(\"pose_bundle\"),\n",
        "#    delete_prefix_on_load=False,  # Use this if downloading is a pain.\n",
        "    zmq_url=zmq_url,\n",
        ")\n",
        "\n",
        "diagram = builder.Build()\n",
        "diagram.set_name(\"pick_and_place\")\n",
        "\n",
        "simulator = Simulator(diagram)\n",
        "station_context = station.GetMyContextFromRoot(simulator.get_mutable_context())\n",
        "# TODO(russt): Add this missing python binding\n",
        "#integrator.set_integral_value(\n",
        "#    integrator.GetMyContextFromRoot(simulator.get_mutable_context()), \n",
        "#        station.GetIiwaPosition(station_context))\n",
        "integrator.GetMyContextFromRoot(simulator.get_mutable_context()).get_mutable_continuous_state_vector().SetFromVector(station.GetIiwaPosition(station_context))\n",
        "\n",
        "simulator.set_target_realtime_rate(1.0)\n",
        "meshcat.start_recording()\n",
        "simulator.AdvanceTo(traj_p_G.end_time() if running_as_notebook else 0.1)\n",
        "meshcat.stop_recording()\n",
        "meshcat.publish_recording()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AeKMhIacWZ00",
        "colab_type": "text"
      },
      "source": [
        "That's quite some diagram we've assembled.  Do you understand all of the parts?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZWXemy-bUAab",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "SVG(pydot.graph_from_dot_data(diagram.GetGraphvizString(max_depth=1))[0].create_svg())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JX1Oq7gfWsBD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}